package com.jtzc.aikf.config;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class EmbeddingStoreConfig {

//    @Autowired
//    private EmbeddingModel embeddingModel;
    @Value("${chroma.url}")
    private String chromaUrl;

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        //创建向量存储
        EmbeddingStore<TextSegment> embeddingStore = ChromaEmbeddingStore.builder()
                .baseUrl(chromaUrl)
                .collectionName("xiaoju-index")
                .logRequests(true)
                .logResponses(true)
                .build();
        return embeddingStore;
    }

//    @Bean
//    public EmbeddingStore<TextSegment> embeddingStore() {
//        //创建向量存储
//        EmbeddingStore<TextSegment> embeddingStore = PineconeEmbeddingStore.builder()
//                .apiKey(System.getenv("PINECONE_API_KEY"))
//                .index("xiaoju-index")//如果指定的索引不存在，将创建一个新的索引
//                .nameSpace("xiaoju-namespace") //如果指定的名称空间不存在，将创建一个新的名称空间
//                .createIndex(PineconeServerlessIndexConfig.builder()
//                        .cloud("AWS") //指定索引部署在 AWS 云服务上。
//                        .region("us-east-1") //指定索引所在的 AWS 区域为 us-east-1。
//                        .dimension(embeddingModel.dimension()) //指定索引的向量维度，该维度与 embeddedModel 生成的向量维度相同。
//                        .build())
//                .build();
//
//        return embeddingStore;
//    }


//    public void test() {
//        try (ChromaDB chroma = new ChromaDBContainer("chromadb/chroma:0.5.4")) {
//            chroma.start();
//
//            EmbeddingStore<TextSegment> embeddingStore = ChromaEmbeddingStore
//                    .builder()
//                    .baseUrl(chroma.getEndpoint())
//                    .collectionName(randomUUID())
//                    .logRequests(true)
//                    .logResponses(true)
//                    .build();
//
//            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
//
//            TextSegment segment1 = TextSegment.from("I like football.");
//            Embedding embedding1 = embeddingModel.embed(segment1).content();
//            embeddingStore.add(embedding1, segment1);
//
//            TextSegment segment2 = TextSegment.from("The weather is good today.");
//            Embedding embedding2 = embeddingModel.embed(segment2).content();
//            embeddingStore.add(embedding2, segment2);
//
//            Embedding queryEmbedding = embeddingModel.embed("What is your favourite sport?").content();
//            EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
//                    .queryEmbedding(queryEmbedding)
//                    .maxResults(1)
//                    .build();
//            List<EmbeddingMatch<TextSegment>> matches = embeddingStore.search(embeddingSearchRequest).matches();
//            EmbeddingMatch<TextSegment> embeddingMatch = matches.get(0);
//
//            System.out.println(embeddingMatch.score()); // 0.8144288493114709
//            System.out.println(embeddingMatch.embedded().text()); // I like football.
//
//            chroma.stop();
//        }
//    }
}
